import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import torch.cuda
import numpy as np
from 封装ST_GCN.model.st_gcn import STGCN


class STGCNTrainer:
    def __init__(self, train_data_path, train_label_path, num_nodes, use_residual, batch_size=32, num_epochs=300, lr=0.001,
                 patience=30,
                 graph_args={'layout': 'openpose','strategy': 'adaptive'}, edge_importance_weighting=True):
        # 检查是否有可用的CUDA（NVIDIA GPU）设备，如果有则使用GPU，否则使用CPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")

        # 从指定路径加载训练数据和对应的标签
        self.train_data = np.load(train_data_path)
        self.train_label = np.load(train_label_path)

        # 将加载的NumPy数组数据转换为PyTorch张量，并将其移动到指定设备（GPU或CPU）上
        self.train_data = torch.tensor(self.train_data, dtype=torch.float32).to(self.device)
        self.train_label = torch.tensor(self.train_label, dtype=torch.long).to(self.device)

        # 创建一个TensorDataset对象，将训练数据和标签组合在一起，方便后续的数据加载
        train_dataset = TensorDataset(self.train_data, self.train_label)
        # 创建一个DataLoader对象，用于批量加载训练数据，设置了批量大小和数据打乱的参数
        self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

        # 计算训练数据中标签的唯一类别数量，这将作为模型的输出类别数
        self.num_classes = len(np.unique(self.train_label.cpu().numpy()))
        # 获取训练数据的通道数，作为模型的输入通道数
        self.in_channels = self.train_data.shape[1]
        # 保存图结构相关的参数，如布局和策略
        self.graph_args = graph_args
        # 保存是否对边进行重要性加权的标志
        self.edge_importance_weighting = edge_importance_weighting
        # 保存节点数量
        self.num_nodes = num_nodes
        # 保存是否使用残差连接的标志
        self.use_residual = use_residual

        # 创建STGCN模型实例，并将其移动到指定设备（GPU或CPU）上
        self.model = STGCN(self.in_channels, self.num_classes, self.graph_args,
                           self.num_nodes,
                           self.edge_importance_weighting,
                           use_residual=self.use_residual).to(self.device)

        # 定义交叉熵损失函数，用于分类任务
        self.criterion = nn.CrossEntropyLoss()
        # 定义Adam优化器，用于更新模型的参数，设置了学习率
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr)

        # 早停策略的相关参数
        # 早停的耐心值，即多少个epoch损失没有改善就停止训练
        self.patience = patience
        # 记录目前损失最小的epoch
        self.best_epoch = 0
        # 记录损失没有改善的epoch数量
        self.counter = 0
        # 初始化最佳训练损失为正无穷大
        self.best_train_loss = float('inf')
        # 训练的总epoch数
        self.num_epochs = num_epochs

    def calculate_accuracy(self, data_loader):
        # 将模型设置为评估模式，关闭一些在训练时使用的特殊层（如Dropout）
        self.model.eval()
        # 初始化正确预测的样本数
        correct = 0
        # 初始化总的样本数
        total = 0
        # 在不计算梯度的上下文中进行操作，以节省内存和计算资源
        with torch.no_grad():
            # 遍历数据加载器中的每个批次
            for inputs, labels in data_loader:
                # 使用模型进行预测
                outputs = self.model(inputs)
                # 获取预测结果中概率最大的类别索引
                _, predicted = torch.max(outputs.data, 1)
                # 累加当前批次的样本数
                total += labels.size(0)
                # 累加当前批次中预测正确的样本数
                correct += (predicted == labels).sum().item()
        # 计算准确率
        accuracy = correct / total
        return accuracy

    def train(self):
        # 开始训练循环，循环次数为预设的epoch数
        for epoch in range(self.num_epochs):
            # 将模型设置为训练模式，开启一些在训练时使用的特殊层（如Dropout）
            self.model.train()
            # 初始化每个epoch的累积损失
            running_loss = 0.0
            # 遍历训练数据加载器中的每个批次
            for i, (inputs, labels) in enumerate(self.train_loader):
                # 清零优化器中的梯度信息，避免梯度累积
                self.optimizer.zero_grad()
                # 使用模型对输入数据进行预测
                outputs = self.model(inputs)
                # 计算预测结果与真实标签之间的损失
                loss = self.criterion(outputs, labels)
                # 反向传播，计算梯度
                loss.backward()
                # 使用优化器更新模型的参数
                self.optimizer.step()
                # 累加当前批次的损失
                running_loss += loss.item()
            # 计算当前epoch的平均训练损失
            train_loss = running_loss / len(self.train_loader)

            # 计算当前epoch的训练准确率
            train_accuracy = self.calculate_accuracy(self.train_loader)
            # 打印当前epoch的训练信息，包括epoch编号、损失和准确率
            print(f'Epoch {epoch + 1}/{self.num_epochs}, Loss: {train_loss}, Train Accuracy: {train_accuracy * 100:.2f}%')

            # 如果当前的训练损失小于之前记录的最佳训练损失
            if train_loss < self.best_train_loss:
                # 更新最佳训练损失
                self.best_train_loss = train_loss
                # 更新最佳epoch编号
                self.best_epoch = epoch
                # 重置损失没有改善的计数器
                self.counter = 0
                print('Saving')
                # 保存当前模型的参数到文件中
                torch.save(self.model.state_dict(),'stgcn_model.pth')
            else:
                # 如果当前损失没有改善，计数器加1
                self.counter += 1
                # 如果计数器达到了早停的耐心值
                if self.counter >= self.patience:
                    print(f'Early stopping at epoch {epoch + 1}. Best epoch was {self.best_epoch + 1}.')
                    # 提前停止训练
                    break